// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/connectors/connector.h>
#include <pollux/exec/cursor.h>
#include <pollux/exec/exchange.h>
#include <pollux/runner/multi_fragment_plan.h>
#include <pollux/runner/runner.h>

namespace kumo::pollux::runner {
    /// Testing proxy for a split source managed by a system with full metadata
    /// access.
    class SimpleSplitSource : public SplitSource {
    public:
        explicit SimpleSplitSource(
            std::vector<std::shared_ptr<connector::ConnectorSplit> > splits)
            : splits_(std::move(splits)) {
        }

        virtual std::vector<SplitAndGroup> getSplits(uint64_t targetBytes) override;

    private:
        std::vector<std::shared_ptr<connector::ConnectorSplit> > splits_;
        int32_t splitIdx_{0};
    };

    /// Testing proxy for a split source factory that uses connector metadata to
    /// enumerate splits. This takes a precomputed split list for each scan.
    class SimpleSplitSourceFactory : public SplitSourceFactory {
    public:
        explicit SimpleSplitSourceFactory(
            melon::F14FastMap<
                core::PlanNodeId,
                std::vector<std::shared_ptr<connector::ConnectorSplit> > > nodeSplitMap)
            : nodeSplitMap_(std::move(nodeSplitMap)) {
        }

        std::shared_ptr<SplitSource> splitSourceForScan(
            const core::TableScanNode &scan) override;

    private:
        melon::F14FastMap<
            core::PlanNodeId,
            std::vector<std::shared_ptr<connector::ConnectorSplit> > >
        nodeSplitMap_;
    };

    /// Runner for in-process execution of a distributed plan.
    class LocalRunner : public Runner,
                        public std::enable_shared_from_this<LocalRunner> {
    public:
        LocalRunner(
            MultiFragmentPlanPtr plan,
            std::shared_ptr<core::QueryCtx> queryCtx,
            std::shared_ptr<SplitSourceFactory> splitSourceFactory)
            : plan_(std::move(plan)),
              fragments_(plan_->fragments()),
              options_(plan_->options()),
              splitSourceFactory_(std::move(splitSourceFactory)) {
            params_.queryCtx = std::move(queryCtx);
        }

        RowVectorPtr next() override;

        std::vector<exec::TaskStats> stats() const override;

        void abort() override;

        void waitForCompletion(int32_t maxWaitMicros) override;

        State state() const override {
            return state_;
        }

    private:
        void start();

        // Creates all stages except for the single worker final consumer stage.
        std::vector<std::shared_ptr<exec::RemoteConnectorSplit> > makeStages();

        std::shared_ptr<SplitSource> splitSourceForScan(
            const core::TableScanNode &scan);

        // Serializes 'cursor_' and 'error_'.
        mutable std::mutex mutex_;

        const MultiFragmentPlanPtr plan_;
        const std::vector<ExecutableFragment> fragments_;
        const MultiFragmentPlan::Options &options_;

        exec::CursorParameters params_;

        tsan_atomic<State> state_{State::kInitialized};

        std::unique_ptr<exec::TaskCursor> cursor_;
        std::vector<std::vector<std::shared_ptr<exec::Task> > > stages_;
        std::exception_ptr error_;
        std::shared_ptr<SplitSourceFactory> splitSourceFactory_;
    };
} // namespace kumo::pollux::runner
